package com.gitee.feizns.explore.data_structure.tree.cus.rbt;

import java.util.Collection;
import java.util.Set;

/**
 * @author feizns
 * @since 2020/1/17
 */
public class RbtImpl4<K extends Comparable<K>, V> implements Rbt<RbtImpl4.Node, K, V> {

    public static void main(String[] args) {
        RbtImpl4 r = new RbtImpl4();
        r.put(1, "1");
        r.put(2, "1");
        r.put(3, "1");
        r.put(4, "1");
        r.put(5, "1");
        r.put(6, "1");
        r.put(7, "1");
        r.put(8, "1");

        r.remove(8);
        System.out.println(r.size());
    }

    private Node root;

    private int size;

    @Override
    public Node root() {
        return root;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public V get(Object key) {
        Node node = node(root, (K) key);
        return node != null ? (V) node.val() : null;
    }

    private Node node(Node n, K k) {
        if ( n != null ) {
            final int com = k.compareTo((K) n.k);
            if ( com < 0 )
                return node(n.left(), k);
            else if ( com > 0 )
                return node(n.right(), k);
            else
                return n;
        }
        return null;
    }

    @Override
    public V put(K key, V value) {
        NAndV put = put(root, key, value);
        root = put.n;
        toB(put.n);
        return put.v;
    }

    private class NAndV {
        Node n;
        V v;
        public NAndV(Node n, V v) {
            this.n = n;
            this.v = v;
        }
    }

    private NAndV put(Node n, K k, V v) {
        if ( n == null ) {
            size++;
            return new NAndV(new Node(k, v), null);
        }

        V ret;
        int comp = k.compareTo((K) n.k);
        if ( comp < 0 ) {
            NAndV put = put(n.left, k, v);
            n.left = put.n;
            ret = v;
        } else if ( comp > 0 ) {
            NAndV put = put(n.right, k, v);
            n.right = put.n;
            ret = v;
        } else {
            ret = (V) n.v;
            n.v = v;
        }

        return new NAndV(balance(n), ret);
    }

    @Override
    public V remove(Object key) {
        NAndV andV = remove(root, (K) key);
        root = andV.n;
        toB(root);
        return andV.v;
    }

    private NAndV remove(Node n, K k) {
        if ( n == null )
            return new NAndV(null, null);

        V ret = null;
        int comp = k.compareTo((K) n.k);
        if ( comp < 0 ) {
            n = moveToLeft(n);
            NAndV andV = remove(n.left, k);
            n.left = andV.n;
            ret = andV.v;
        } else if ( comp > 0 ) {
            n = moveToRight(n);
            NAndV andV = remove(n.right, k);
            n.right = andV.n;
            ret = andV.v;
        } else {
            if ( n.left == null && n.right == null ) {
                size--;
                return new NAndV(null, (V) n.v);
            } else if ( n.left != null ) {
                Node max = max(n.left);
                n.k = max.k;
                n.v = max.v;
                n = moveToLeft(n);
                NAndV andV = remove(n.left, (K) max.k);
                n.left = andV.n;
            } else {
                Node min = min(n.right);
                n.k = min.k;
                n.v = min.v;
                n = moveToRight(n);
                NAndV andV = remove(n.right, (K) min.k);
                n.right = andV.n;
            }
        }

        return new NAndV(balance(n), ret);
    }

    private static Node moveToLeft(Node n) {
        if ( n.left != null ) {
            if ( !isR(n.left) && !isR(n.left.left) && !isR(n.left.right) ) {
                if ( n.right != null ) {
                    if ( !isR(n.right) && !isR(n.right.left) && !isR(n.right.right) ) {
                        flipColorsAA(n);
                    } else {
                        if ( isR(n.right.left) )
                            n.right = rightRotate(n);
                        n = leftRotate(n);
                        flipColorsAA(n.left);
                    }
                }
            }
        }
        return n;
    }

    private static Node moveToRight(Node n) {
        if ( n.right != null ) {
            if ( !isR(n.right) && !isR(n.right.left) && !isR(n.right.right) ) {
                if ( n.left != null ) {
                    if ( !isR(n.left) && !isR(n.left.left) && !isR(n.left.right) ) {
                        flipColorsAA(n);
                    } else {
                        if ( isR(n.left.right) )
                            n.left = leftRotate(n);
                        n = rightRotate(n);
                        flipColorsAA(n.right);
                    }
                }
            }
        }
        return n;
    }

    private static Node max(Node n) {
        Node tmp = n;
        while ( tmp.right != null )
            tmp = tmp.right;
        return tmp;
    }

    private static Node min(Node n) {
        Node tmp = n;
        while ( tmp.left != null )
            tmp = tmp.left;
        return tmp;
    }

    @Override
    public void clear() {

    }

    @Override
    public Set<K> keySet() {
        return null;
    }

    @Override
    public Collection<V> values() {
        return null;
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        return null;
    }

    public static class EntryImpl<K, V> implements Entry<K, V> {
        K k;
        V v;
        public EntryImpl(K k, V v) {
            this.k = k;
            this.v = v;
        }
        @Override
        public K getKey() {
            return k;
        }
        @Override
        public V getValue() {
            return v;
        }
        @Override
        public V setValue(V value) {
            V old = v;
            this.v = value;
            return old;
        }
    }

    private static Node balance(Node n) {
        if ( !isR(n) && isR(n.right) )
            n = leftRotate(n);
        if ( isR(n.left) && isR(n.left.left) )
            n = rightRotate(n);
        if ( !isR(n) && isR(n.left) && isR(n.right) )
            flipColors(n);
        return n;
    }

    private static void flipColors(Node n) {
        toR(n);
        toB(n.left);
        toB(n.right);
    }

    private static void flipColorsAA(Node n) {
        toB(n);
        toR(n.left);
        toR(n.right);
    }

    private static Node toR(Node n) {
        if ( n != null )
            n.color = RED;
        return n;
    }

    private static Node toB(Node n) {
        if ( n != null )
            n.color = BLACK;
        return n;
    }

    private static Node leftRotate(Node n) {
        Node right = n.right;
        n.right = right.left;
        right.left = n;
        right.color = n.color;
        n.color = RED;
        return right;
    }

    private static Node rightRotate(Node n) {
        Node left = n.left;
        n.left = left.right;
        left.right = n;
        left.color = n.color;
        n.color = RED;
        return left;
    }

    private static boolean isR(Node n) {
        return n != null && n.color == RED;
    }

    private static boolean BLACK = true;
    private static boolean RED = false;

    static class Node<K extends Comparable<K>, V> implements com.gitee.feizns.explore.data_structure.tree.cus.Node<Node, K, V> {
        boolean color;
        K k;
        V v;
        Node left, right;
        public Node(K k, V v) {
            this.k = k;
            this.v = v;
        }
        public Node(K k, V v, boolean color) {
            this(k, v);
            this.color = color;
        }
        @Override
        public K key() {
            return k;
        }
        @Override
        public void key(K k) {
            this.k = k;
        }
        @Override
        public V val() {
            return v;
        }
        @Override
        public void val(V v) {
            this.v = v;
        }
        @Override
        public Node left() {
            return left;
        }
        @Override
        public void left(Node left) {
            this.left = left;
        }
        @Override
        public Node right() {
            return right;
        }
        @Override
        public void right(Node right) {
            this.right = right;
        }
    }

}
